#include <QJsonDocument>
#include <QJsonObject>
#include "model.h"
#include "mainwindow.h"
#include "serverssettingsdialog.h"

Model::Model(QObject* parent) : QObject(parent){}

Model::~Model() {
    this->freeProcess();
    this->freeTcpClient();
    this->freeEventLoop();
}

void Model::setModelInfo(
    const QString& modelNameOrServerURL, const QString& jsonTemplate, const QString& keys
) {
    this->currentModel = modelNameOrServerURL;
    this->jsonTemplate = jsonTemplate;
    this->keys = keys;
}

void Model::setQuestion(const QString& question, const QString& recordStr) {
    this->originalQuestion = question;
    if (recordStr.length() == 0)
        this->question = "User: " + question;
    else
        this->question = recordStr + "User: " + question;
}

void Model::initProcess() {
    if (this->process != nullptr)
        this->freeProcess();
    this->process = new QProcess(this);
    connect(this->process, &QProcess::readyReadStandardOutput, this, &Model::slot_readFromQProcess);
}

void Model::freeProcess() {
    if (this->process != nullptr) {
        this->process->close();
        delete this->process;
        this->process = nullptr;
        if (!QThread::currentThread()->isInterruptionRequested())
            QThread::currentThread()->requestInterruption();
    }
}

bool Model::initTcpClient() {
    if (this->tcpClient != nullptr)
        this->freeTcpClient();
    this->tcpClient = new TcpClient(this, this->currentModel);
    if (!this->tcpClient->initSocket())
        return false;
    connect(this->tcpClient, &TcpClient::signal_response, this, &Model::slot_readFromTcpClient);
    connect(this->tcpClient, &TcpClient::signal_finished, this, &Model::slot_readFinalFromTcpClient);
    return true;
}

void Model::freeTcpClient() {
    if (this->tcpClient != nullptr) {
        this->tcpClient->freeSocket();
        delete this->tcpClient;
        this->tcpClient = nullptr;
        if (!QThread::currentThread()->isInterruptionRequested())
            QThread::currentThread()->requestInterruption();
    }
}

void Model::initEventLoop() {
    this->freeEventLoop();
    this->loop = new QEventLoop(this);
}

void Model::freeEventLoop() {
    if (this->loop != nullptr) {
        this->loop->quit();
        delete this->loop;
        this->loop = nullptr;
        if (!QThread::currentThread()->isInterruptionRequested())
            QThread::currentThread()->requestInterruption();
    }
}

QString Model::pickValues(const QByteArray& json) {
    int startIdx = json.indexOf("{");
    if (startIdx == -1)
        return "";
    int endIdx = json.lastIndexOf("}");
    if (endIdx == -1)
        return "";
    QJsonDocument jsonDoc = QJsonDocument::fromJson(
        json.mid(startIdx, endIdx + 1 - startIdx)
    );
    if (jsonDoc.isNull())
        return "";
    if (this->keys.length() != 0) {
        QJsonObject jsonObj = jsonDoc.object();
        QString res = "";
        QStringList keyList = this->keys.split(",");
        int maxIndex = keyList.length() - 1;
        for (int i = 0; i < maxIndex; ++ i)
            res += jsonObj[keyList[i]].toString() + "_";
        return res + jsonObj[keyList[maxIndex]].toString();
    }
    return jsonDoc.toJson(QJsonDocument::Compact);
}

void Model::work() {
    this->answer.clear();
    emit signal_updateOutputArea(this->originalQuestion, this->answer);
    if (this->jsonTemplate.length()) {
        QString postData = ServersSettingsDialog::getFinalJson(
            this->jsonTemplate, this->question
        );
        if (!this->initTcpClient()) {
            this->freeTcpClient();
            emit signal_finished(this->originalQuestion, tr("Failed to create client!"));
            return;
        }
        if (!this->tcpClient->postRequest(postData.toUtf8())) {
            this->freeTcpClient();
            emit signal_finished(this->originalQuestion, tr("Failed to post request!"));
            return;
        }
        if (this->loop != nullptr)
            delete this->loop;
        this->initEventLoop();
        this->loop->exec();
        emit signal_finished(this->originalQuestion, this->answer);
        return;
    }
    this->initProcess();
    
    QStringList args = { "run", this->currentModel };
    this->process->start("ollama", args);
    this->process->waitForStarted();
    
    this->process->write(this->question.toUtf8() + "\n");
    this->process->closeWriteChannel();
    this->process->waitForFinished(60 * 60 * 1000);
    this->freeProcess();
    emit signal_finished(this->originalQuestion, this->answer);
}

void Model::slot_readFromQProcess() {
    if (QThread::currentThread()->isInterruptionRequested()) {
        this->process->close();
        return;
    }
    QByteArray output = this->process->readAllStandardOutput();
    QString partialAnswer = QString(output);
    this->answer += partialAnswer;
    emit signal_updateOutputArea(this->originalQuestion, this->answer);
}

void Model::slot_readFromTcpClient(const QByteArray& partialResponse) {
    if (QThread::currentThread()->isInterruptionRequested()) {
        this->freeEventLoop();
        return;
    }
    QString partialAnswer = this->pickValues(partialResponse);
    if (partialAnswer.length() == 0) {
        bool canTry = (this->currentJsonContent.length() != 0);
        this->currentJsonContent += partialResponse;
        if (canTry) {
            partialAnswer = this->pickValues(this->currentJsonContent);
            if (partialAnswer.length() == 0)
                return;
            this->currentJsonContent.clear();
        } else {
            return;
        }
    }
    this->answer += partialAnswer;
    emit signal_updateOutputArea(this->originalQuestion, this->answer);
}

void Model::slot_readFinalFromTcpClient(const QByteArray& partialResponse) {
    if (QThread::currentThread()->isInterruptionRequested()) {
        this->freeEventLoop();
        return;
    }
    if (partialResponse.length()) {
        QString partialAnswer = this->pickValues(partialResponse);
        if (partialAnswer.length() == 0) {
            bool canTry = (this->currentJsonContent.length() != 0);
            this->currentJsonContent += partialResponse;
            if (canTry)
                partialAnswer = this->pickValues(this->currentJsonContent);
            this->currentJsonContent.clear();
        }
        this->answer += partialAnswer;
    }
    this->freeEventLoop();
}